import os
import math
import struct
from datetime import datetime
import backtrader as bt
import numpy as np
import pandas as pd
from Units import QUnits, TdxSecurity
import plotly.express as px
import plotly.graph_objects as go

pd.set_option('expand_frame_repr', False)  # 当列太多时不换行
# 负责读取本地通达信数据
class TDXLoader:
    def __init__(self):
        unit = QUnits()
        # TDX_DIR 为通达信根目录
        self.TDX_DIR = os.path.dirname(unit.XML2ObjFromFile('./setting.xml').setting.tdx.path)

    def resolve(self, *args, **kwargs):
        return os.path.join(self.TDX_DIR, *args, **kwargs)

    # 读取K线源文件
    def _read_kline(self, filepath, pack='lllllfll'):
        with open(filepath, "rb") as f:
            usecols = "trade_date open high low close amount vol openinterest".split()
            buffers = []
            while True:
                buffer = f.read(32)
                if not buffer:
                    break
                buffer = struct.unpack(pack, buffer)
                if(pack == 'lllllfll'):
                    buffers.append(buffer)
                else:
                    buffers.append((self.get_date_str(buffer[0], buffer[1]), buffer[2], buffer[3], buffer[4], buffer[5], buffer[6], buffer[7], buffer[8]))
            kline = pd.DataFrame(buffers, columns=usecols)
        kline["trade_date"] = kline["trade_date"].astype(str)
        price_columns = ["open", "high", "low", "close"]
        kline[price_columns] = kline[price_columns].apply(lambda x: x / 100)
        return kline

    # 根据二进制前两段拿到日期分时
    def get_date_str(self, h1, h2) -> str:  # H1->0,1字节; H2->2,3字节;
        year = math.floor(h1 / 2048) + 2004  # 解析出年
        month = math.floor(h1 % 2048 / 100)  # 月
        day = h1 % 2048 % 100  # 日
        hour = math.floor(h2 / 60)  # 小时
        minute = h2 % 60  # 分钟
        if hour < 10:  # 如果小时小于两位, 补0
            hour = "0" + str(hour)
        if minute < 10:  # 如果分钟小于两位, 补0
            minute = "0" + str(minute)
        return str(year) + "-" + str(month) + "-" + str(day) + " " + str(hour) + ":" + str(minute)

    # 获取日K线数据
    def get_kline_daily(self, ts_code, cycle=None):
        if(cycle in ('1h', '30m', '15m', '5m')):
            filename = ts_code.split(".")[1] + ts_code.split(".")[0] + ".lc5"
            filepath = self.resolve("vipdoc", ts_code.split(".")[1], "fzline", filename)
            pack = 'HHffffllf'
        elif(cycle == '1m'):
            filename = ts_code.split(".")[1] + ts_code.split(".")[0] + ".lc1"
            filepath = self.resolve("vipdoc", ts_code.split(".")[1], "minline", filename)
            pack = 'HHffffllf'
        else:
            filename = ts_code.split(".")[1] + ts_code.split(".")[0] + ".day"
            filepath = self.resolve("vipdoc", ts_code.split(".")[1], "lday", filename)
            pack = 'lllllfll'
        kline = self._read_kline(filepath, pack)
        kline["ts_code"] = ts_code
        kline.index = pd.to_datetime(kline["trade_date"])
        kline.index.name = "index"
        kline = kline.rename(columns={"vol": "volume"})
        usecols = (
            "ts_code trade_date open high low close amount volume openinterest".split()
        )
        return kline[usecols]


# Create a Stratey
class TestStrategy(bt.Strategy):

    def log(self, txt, dt=None):
        ''' 提供记录功能'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):
        # 引用到输入数据的close价格
        self.dataclose = self.datas[0].close

    def next(self):
        # 目前的策略就是简单显示下收盘价。
        self.log('Close, %.2f' % self.dataclose[0])

if __name__ == "__main__":
    loader = TDXLoader()
    stock_df = loader.get_kline_daily("300667.sz", '1d')
    print(stock_df)
    print(type(stock_df))
    cerebro = bt.Cerebro()
    cerebro.addstrategy(TestStrategy)
    start_date = datetime(2020, 9, 30)  # 回测开始时间
    end_date = datetime(2021, 9, 30)  # 回测结束时间
    data = bt.feeds.PandasData(dataname=stock_df, fromdate=start_date, todate=end_date)  # 加载数据
    cerebro.adddata(data)  # 将数据传入回测系统
    cerebro.broker.setcash(100000.0)
    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())
    cerebro.run()
    print('Final Portfolio Value: %.2f' % cerebro.broker.getvalue())
    cerebro.plot()
    '''dKfig = go.Figure(data=go.Candlestick(
        x=dKline['trade_date'],
        open=dKline['open'],  # 字段数据必须是元组、列表、numpy数组、或者pandas的Series数据
        high=dKline['high'],
        low=dKline['low'],
        close=dKline['close']
    ))
    dKfig.update(layout_xaxis_rangeslider_visible=False)
    dKfig.show()'''